import numpy as np
import matplotlib as MP
import pylab as P
#--------------------
def rgrid():
    import grid_data as GD 
    lfluid,lforcing,lparticle,limmersed_boundary,n1,n2,length,xx0,yy0 = GD.rgrid_param();
    pgrid(n1,n2,length,xx0,yy0)
    if (limmersed_boundary == 1) :
        rbound=np.loadtxt('immersed_boundary_points.ascii')
        P.plot(rbound[:,0],rbound[:,1],'r.')
# --------------------------------------
def pgrid(n1=64,n2=64,length=2*np.pi,xx0=-np.pi,yy0=-np.pi):
    dx=length/n1
    dy=length/n2
    xx=xx0+np.linspace(0,n1,num=n1)*dx
    yy=yy0+np.linspace(0,n2,num=n2)*dy
    P.clf()
    for ix in range(0,n1):
        xvert=xx[ix]+np.zeros(n2)
        P.plot(xvert,yy,'b-')
    for iy in range(0,n2):
        yhoriz=yy[iy]+np.zeros(n1)
        P.plot(xx,yhoriz,'b-')
#************************************#
def vel(ptype='contour',nevery=5):
    import grid_data as GD
    lfluid,lforcing,lparticle,limmersed_boundary,n1,n2,length,xx0,yy0 = GD.rgrid_param();
    #pgrid(n1,n2,length,xx0,yy0)
    if (limmersed_boundary == 1) :
        rbound=np.loadtxt('immersed_boundary_points.ascii')
        P.plot(rbound[:,0],rbound[:,1],'r.')
    dx=length/n1
    dy=length/n2
    xx=xx0+np.linspace(0,n1,num=n1)*dx
    yy=yy0+np.linspace(0,n2,num=n2)*dy
    ux=np.loadtxt('data/ux.ascii')
    uy=np.loadtxt('data/uy.ascii')
    if (ptype=='contour'):
        P.contourf(ux)
    if (ptype=='quiver'):
        P.quiver(xx[::nevery],yy[::nevery],ux[::nevery],uy[::nevery],units='width')
#************************************#
def omega(ptype='contour',nevery=5,fname='data/omega.ascii'):
    import grid_data as GD
    lfluid,lforcing,lparticle,limmersed_boundary,n1,n2,length,xx0,yy0 = GD.rgrid_param();
    #pgrid(n1,n2,length,xx0,yy0)
    if (limmersed_boundary == 1) :
        rbound=np.loadtxt('immersed_boundary_points.ascii')
        P.plot(rbound[:,0],rbound[:,1],'r.')
    dx=length/n1
    dy=length/n2
    xx=xx0+np.linspace(0,n1,num=n1)*dx
    yy=yy0+np.linspace(0,n2,num=n2)*dy
    omega=np.loadtxt(fname)
    if (ptype=='contour'):
        P.contourf(omega)
#************************************#
def stream(nevery=5,fname='data/stream.ascii',nlines=64,lfilled=True):
    import grid_data as GD
    lfluid,lforcing,lparticle,limmersed_boundary,n1,n2,length,xx0,yy0 = GD.rgrid_param();
    #pgrid(n1,n2,length,xx0,yy0)
    if (limmersed_boundary == 1) :
        rbound=np.loadtxt('immersed_boundary_points.ascii')
        P.plot(rbound[:,0],rbound[:,1],'r.')
    dx=length/n1
    dy=length/n2
    xx=xx0+np.linspace(0,n1,num=n1)*dx
    yy=yy0+np.linspace(0,n2,num=n2)*dy
    X, Y = np.meshgrid(xx, yy)	
    stream=np.loadtxt(fname)
    if (lfilled) :
      P.contourf(X,Y,stream,nlines)
    else :
      P.contour(X,Y,stream,nlines)
#************************************#
def pforce():
    ffile=np.loadtxt('force.ascii')
    P.plot(ffile[:,0],ffile[:,1])
    P.plot(ffile[:,0],ffile[:,2])
#************************************#
if __name__ == '__main__':
  main();
else:
  print 'Importing Vizualizer for spectral soap'
